/*
** Xin YUAN, 2023, BSD (2)
*/

////////////////////////////////////////////////////////////////////////////////

#include "base/_w_thread_pool.h"

////////////////////////////////////////////////////////////////////////////////

inline uint32_t get_processor_number() noexcept
{
	SYSTEM_INFO si;
	::GetSystemInfo(&si);
	return si.dwNumberOfProcessors;
}

//handle_server_data_base

class handle_server_data_base : public handle_data_base
{
private:
	//io_data
	/*! \brief A structure for basic overlapped.

	A structure for basic overlapped.
	*/
	struct io_data
	{
		WSAOVERLAPPED overlapped;   //!< The basic overlapped member.
		void* pv{nullptr};
		int op_type{OP_TYPE_NONE};  //!< The operation type. See OP_TYPE_*.
	};

	io_data idRead;
	io_data idWrite;
	PTP_IO pIO;
	PTP_TIMER pTimer;

	volatile LONG lCloseFlag;
	volatile LONG lCallFlag;

protected:
	handle_server_data_base() noexcept : pIO(NULL), pTimer(NULL), lCloseFlag(0), lCallFlag(0)
	{
	}
	handle_server_data_base(const handle_server_data_base&) = delete;
	~handle_server_data_base() noexcept
	{
	}
	handle_server_data_base& operator=(const handle_server_data_base&) = delete;

private:
	friend class server_base;
};

//server_base

class server_base : public _w_thread_pool_base
{
private:
	typedef _w_thread_pool_base  baseClass;

protected:
	typedef server_util  utilClass;

protected:
	server_base() noexcept
	{
	}
	server_base(const server_base&) = delete;
	~server_base() noexcept
	{
	}

	server_base& operator=(const server_base&) = delete;

protected:
	bool init(utilClass::server_callback cb, uint32_t tm) noexcept
	{
		//callback
		m_util.callback = cb;
		//timeout
		m_util.init_timeout(tm);
		if ( !baseClass::init() )
			return false;
		return true;
	}

public:
	//Do not reuse handle_server_data_base in used list after calling this method
	//  Socket close will not be graceful after this method calling.
	void Shutdown() noexcept
	{
		baseClass::dump();
	}

	bool StartListen(handle_server_data_base* p, uint32_t port) noexcept
	{
		if ( ::InterlockedCompareExchange(&m_lExit, 1, 1) )
			return true;
		//addr
		SOCKADDR_STORAGE ss;
		int len;
		if ( !_w_socket_helper::FillSockAddr(ss, len, AF_INET, port) )
			return false;
		//socket
		ns_socket sock;
		if ( !sock.GetSocket().Create(AF_INET, SOCK_STREAM, IPPROTO_TCP) )
			return false;
		if ( !sock.GetSocket().Bind((SOCKADDR*)&ss, len) )
			return false;
		if ( !sock.GetSocket().Listen() )
			return false;
		//IOCP
		PTP_IO pIO = ::CreateThreadpoolIo((HANDLE)(SOCKET)(sock.GetSocket()),
						&io_completion_callback, this,
						&m_cbe);
		if ( pIO == NULL ) //::GetLastError()
			return false;
		//work item
		p->sock = rv_forward(sock);
		p->pIO = pIO;
		p->pTimer = NULL;
		p->op_type = OP_TYPE_NONE;
		p->lCloseFlag = 0;
		p->index = (uintptr_t)this;
		if ( !::TrySubmitThreadpoolCallback(&simple_callback, p, &m_cbe) ) { //::GetLastError()
			p->sock.GetSocket().Close();
			::CloseThreadpoolIo(p->pIO);
			return false;
		}
		return true;
	}
	bool StartConnect(handle_server_data_base* p, const char_s* host, uint32_t port) noexcept
	{
		if ( ::InterlockedCompareExchange(&m_lExit, 1, 1) )
			return true;
		//addr
		SOCKADDR_STORAGE ss;
		int len;
		if ( !host_to_sockaddr(host, port, ss, len) )
			return false;
		//socket
		ns_socket sock;
		if ( !sock.GetSocket().Create(AF_INET, SOCK_STREAM, IPPROTO_TCP) )
			return false;
		{
			//bind local address
			SOCKADDR_STORAGE local_ss;
			int local_len;
			if ( !_w_socket_helper::FillSockAddr(local_ss, local_len, AF_INET, 0) )
				return false;
			if ( !sock.GetSocket().Bind((SOCKADDR*)&local_ss, local_len) )
				return false;
		} //end block
		//function
		LPFN_CONNECTEX fnConnectEx = nullptr;
		GUID guidConnectEx = WSAID_CONNECTEX;
		DWORD dwBytes = 0;
		if ( !sock.GetSocket().Ioctl(SIO_GET_EXTENSION_FUNCTION_POINTER,
						&guidConnectEx, sizeof(GUID),
						&fnConnectEx, sizeof(LPFN_CONNECTEX),
						&dwBytes, NULL) )
			return false;
		//IOCP
		PTP_IO pIO = ::CreateThreadpoolIo((HANDLE)(SOCKET)(sock.GetSocket()),
						&io_completion_callback, this,
						&m_cbe);
		if ( pIO == NULL ) //::GetLastError()
			return false;
		//timer
		PTP_TIMER pTimer = ::CreateThreadpoolTimer(&timer_callback, p, &m_cbe);
		if ( pTimer == NULL ) { //::GetLastError()
			::CloseThreadpoolIo(pIO);
			return false;
		}
		//connect
		::ZeroMemory(&(p->idWrite.overlapped), sizeof(WSAOVERLAPPED));
		p->idWrite.op_type = OP_TYPE_CONNECT;
		p->idWrite.pv = p;
		p->sock = rv_forward(sock);
		p->pIO = pIO;
		p->pTimer = pTimer;
		p->op_type = OP_TYPE_NONE;
		p->lCloseFlag = 0;
		p->send_len = p->act_sent_len = 0;
		p->index = (uintptr_t)this;
		::StartThreadpoolIo(p->pIO);
		if ( !fnConnectEx(p->sock.GetSocket(), (sockaddr*)&ss, len,
					NULL, 0, NULL,
					&(p->idWrite.overlapped)) ) {
			if ( ::WSAGetLastError() != ERROR_IO_PENDING ) {
				p->sock.GetSocket().Close();
				::CancelThreadpoolIo(p->pIO);
				::CloseThreadpoolIo(p->pIO);
				::CloseThreadpoolTimer(p->pTimer);
				return false;
			}
		}
		return true;
	}

protected:
	bool _BeginInput(handle_server_data_base* p, uint32_t uSendLength) noexcept
	{
		assert( uSendLength > 0 && uSendLength <= HANDLE_BUFFER_SIZE );
		//acquire
		while ( ::InterlockedCompareExchange(&(p->lCallFlag), 1, 0) )
			;
		//check
		if ( ::InterlockedCompareExchange(&m_lExit, 1, 1) )
			return true;
		if ( ::InterlockedCompareExchange(&(p->lCloseFlag), 1, 1) )
			return true;
		if ( ::InterlockedCompareExchange((LONG volatile*)&(p->send_len), uSendLength, 0) ) {
			::InterlockedCompareExchange(&(p->lCallFlag), 0, 1);
			return false;
		}
		return true;
	}
	bool _EndInput(handle_server_data_base* p) noexcept
	{
		bool b = start_send(p);
		::InterlockedCompareExchange(&(p->lCallFlag), 0, 1);
		return b;
	}

	//This method must be called in DoSent() using its input parameter p before the next _StartInput() calling.
	static void _FinishSend(handle_server_data_base* p) noexcept
	{
		p->act_sent_len = 0;
		::InterlockedExchange((LONG volatile*)&(p->send_len), 0);
	}

	bool _MarkClose(handle_server_data_base* p) noexcept
	{
		//acquire
		while( ::InterlockedCompareExchange(&(p->lCallFlag), 1, 0) )
			;
		//process
		bool b = true;
		LONG lOrg = ::InterlockedCompareExchange(&(p->lCloseFlag), 1, 0);
		if ( lOrg == 0 ) {
			if ( p->op_type != OP_TYPE_ACCEPT ) {
				if ( !::TrySubmitThreadpoolCallback(&simple_callback, p, &m_cbe) ) //::GetLastError()
					b = false;
			}
		}
		//finish
		::InterlockedCompareExchange(&(p->lCallFlag), 0, 1);
		return b;
	}

private:
//callbacks
	static VOID CALLBACK simple_callback(
					PTP_CALLBACK_INSTANCE Instance,
					PVOID                 Context
					) noexcept
	{
		handle_server_data_base* p = (handle_server_data_base*)Context;
		server_base* pThis = (server_base*)(void*)(p->index);
		if ( ::InterlockedCompareExchange(&pThis->m_lExit, 1, 1) )
			return ;
		if ( !::InterlockedCompareExchange(&(p->lCloseFlag), 1, 1) ) {
			//listen
			uint32_t num = get_processor_number() + 2;
			if ( num > _POOL_MAX_THREAD_NUM / 2 )
				num = _POOL_MAX_THREAD_NUM / 2;
			for ( uint32_t i = num; i > 0; -- i ) {
				handle_server_data_base* pA = (handle_server_data_base*)pThis->m_util.callback(utilClass::CB_BEGIN_ACCEPT, NULL, true, pThis);
				if ( pA == NULL )
					continue;
				if ( !pThis->start_accept(p, pA) )
					pThis->m_util.callback(utilClass::CB_CLOSED_END, pA, true, pThis);
			}
			return ;
		}
		//close
		while ( ::InterlockedCompareExchange(&(p->lCallFlag), 1, 1) )
			;
		close_handle(p);
		pThis->m_util.callback(utilClass::CB_CLOSED_BEGIN, p, true, pThis);
		while ( ::InterlockedCompareExchange(&(p->lCallFlag), 1, 1) )
			;
		pThis->m_util.callback(utilClass::CB_CLOSED_END, p, true, pThis);
	}
	bool start_accept(handle_server_data_base* pL, handle_server_data_base* pA) noexcept
	{
		if ( ::InterlockedCompareExchange(&m_lExit, 1, 1) )
			return true;
		if ( !pA->sock.GetSocket().Create(AF_INET, SOCK_STREAM, IPPROTO_TCP) )
			return false;
		::ZeroMemory(&(pA->idRead.overlapped), sizeof(WSAOVERLAPPED));
		pA->idRead.op_type = OP_TYPE_ACCEPT;
		pA->idRead.pv = pA;
		pA->idWrite.pv = pL;
		pA->op_type = OP_TYPE_ACCEPT;
		pA->lCloseFlag = 0;
		pA->send_len = pA->act_sent_len = 0;
		pA->index = (uintptr_t)this;
		pA->pIO = NULL;
		pA->pTimer = NULL;
		::StartThreadpoolIo(pL->pIO);
		if ( !pL->sock.GetSocket().Accept(pA->sock.GetSocket(), pA->recv_buffer, &(pA->idRead.overlapped)) ) {
			if ( ::WSAGetLastError() != ERROR_IO_PENDING ) {
				::CancelThreadpoolIo(pL->pIO);
				pA->sock.GetSocket().Close();
				return false;
			}
		}
		return true;
	}

	static VOID CALLBACK timer_callback(
					PTP_CALLBACK_INSTANCE Instance,
					PVOID                 Context,
					PTP_TIMER             Timer
					) noexcept
	{
		handle_server_data_base* p = (handle_server_data_base*)Context;
		server_base* pThis = (server_base*)(void*)(p->index);
		if ( ::InterlockedCompareExchange(&pThis->m_lExit, 1, 1) )
			return ;
		pThis->set_timer(p, true);
		pThis->m_util.callback(utilClass::CB_TIMEOUT, p, false, pThis);
		pThis->check_and_enable_timer(p);
	}

	static VOID CALLBACK io_completion_callback(
					PTP_CALLBACK_INSTANCE Instance,
					PVOID                 Context,
					PVOID                 Overlapped,
					ULONG                 IoResult,
					ULONG_PTR             NumberOfBytesTransferred,
					PTP_IO                Io
					) noexcept
	{
		server_base* pThis = (server_base*)Context;
		if ( ::InterlockedCompareExchange(&pThis->m_lExit, 1, 1) )
			return ;
		handle_server_data_base::io_data* pD = (handle_server_data_base::io_data*)Overlapped;
		handle_server_data_base* p = (handle_server_data_base*)(pD->pv);
		if ( ::InterlockedCompareExchange(&(p->lCloseFlag), 1, 1) )
			return ;
		//accept
		if ( pD->op_type == OP_TYPE_ACCEPT ) {
			pThis->process_accept(IoResult, (handle_server_data_base*)(p->idWrite.pv), p);
			return ;
		}
		//connect & r/w
		switch ( pD->op_type ) {
		case OP_TYPE_CONNECT:
			pThis->process_connect(IoResult, p);
			break;
		case OP_TYPE_RECV:
			pThis->process_recv(IoResult, NumberOfBytesTransferred, p);
			break;
		case OP_TYPE_SEND:
			pThis->process_sent(IoResult, NumberOfBytesTransferred, p);
			break;
		default:
			break;
		}
		pThis->check_and_enable_timer(p);
	}

	void process_accept(ULONG uIoResult, handle_server_data_base* pL, handle_server_data_base* pA) noexcept
	{
		//pA
		pA->op_type = OP_TYPE_NONE;
		bool b = true;
		if ( uIoResult != NO_ERROR ) {
			b = false;
		}
		else {
			pA->pIO = ::CreateThreadpoolIo((HANDLE)(SOCKET)(pA->sock.GetSocket()),
							&io_completion_callback, this,
							&m_cbe);
			if ( pA->pIO == NULL ) { //::GetLastError()
				b = false;
			}
			else {
				pA->pTimer = ::CreateThreadpoolTimer(&timer_callback, pA, &m_cbe);
				if ( pA->pTimer == NULL ) { //::GetLastError()
					b = false;
				}
				else {
					if ( !start_recv(pA) )
						b = false;
				}
			}
		}
		if ( !b )
			pA->op_type = OP_TYPE_ACCEPT;
		m_util.callback(utilClass::CB_ACCEPTED, pA, b, this);
		if ( !b ) {
			//close
			close_handle(pA);
			m_util.callback(utilClass::CB_CLOSED_BEGIN, pA, true, this);
			m_util.callback(utilClass::CB_CLOSED_END, pA, true, this);
		}
		//new accept
		pA = (handle_server_data_base*)m_util.callback(utilClass::CB_BEGIN_ACCEPT, NULL, true, this);
		if ( pA == NULL )
			return ;
		if ( !start_accept(pL, pA) ) {
			m_util.callback(utilClass::CB_CLOSED_END, pA, true, this);
			return ;
		}
	}
	void process_connect(ULONG uIoResult, handle_server_data_base* p) noexcept
	{
		bool b = true;
		if ( uIoResult != NO_ERROR ) {
			b = false;
		}
		else {
			if ( !start_recv(p) )
				b = false;
		}
		m_util.callback(utilClass::CB_CONNECTED, p, b, this);
	}
	void process_recv(ULONG uIoResult, ULONG_PTR NumberOfBytesTransferred, handle_server_data_base* p) noexcept
	{
		set_timer(p, true);
		if ( uIoResult != NO_ERROR ) {
			m_util.callback(utilClass::CB_RECV, p, false, this);
			return ;
		}
		//recv
		p->recv_len = (uint32_t)NumberOfBytesTransferred;
		bool bZero = p->recv_len == 0;
		m_util.callback(utilClass::CB_RECV, p, true, this);
		//close
		if ( ::InterlockedCompareExchange(&(p->lCloseFlag), 1, 1) )
			return ;
		//next
		if ( !bZero ) {
			if ( !start_recv(p) )
				m_util.callback(utilClass::CB_RECV, p, false, this);
			return ;
		}
	}
	void process_sent(ULONG uIoResult, ULONG_PTR NumberOfBytesTransferred, handle_server_data_base* p) noexcept
	{
		set_timer(p, true);
		if ( uIoResult != NO_ERROR ) {
			m_util.callback(utilClass::CB_SENT, p, false, this);
			return ;
		}
		//sent
		p->act_sent_len += (uint32_t)NumberOfBytesTransferred;
		if ( p->act_sent_len >= p->send_len ) {
			m_util.callback(utilClass::CB_SENT, p, true, this);
			return ;
		}
		//next
		if ( !start_send(p) ) {
			m_util.callback(utilClass::CB_SENT, p, false, this);
			return ;
		}
	}

	bool start_recv(handle_server_data_base* p) noexcept
	{
		if ( ::InterlockedCompareExchange(&m_lExit, 1, 1) )
			return true;
		if ( ::InterlockedCompareExchange(&(p->lCloseFlag), 1, 1) )
			return true;
		::ZeroMemory(&(p->idRead.overlapped), sizeof(WSAOVERLAPPED));
		p->idRead.op_type = OP_TYPE_RECV;
		p->idRead.pv = p;
		::StartThreadpoolIo(p->pIO);
		if ( !p->sock.GetSocket().Read(p->recv_buffer, HANDLE_BUFFER_SIZE, NULL,
						&(p->idRead.overlapped)) ) {
			if ( ::WSAGetLastError() != WSA_IO_PENDING ) {
				::CancelThreadpoolIo(p->pIO);
				return false;
			}
		}
		set_timer(p, false);
		return true;
	}
	bool start_send(handle_server_data_base* p) noexcept
	{
		if ( ::InterlockedCompareExchange(&m_lExit, 1, 1) )
			return true;
		if ( ::InterlockedCompareExchange(&(p->lCloseFlag), 1, 1) )
			return true;
		::ZeroMemory(&(p->idWrite.overlapped), sizeof(WSAOVERLAPPED));
		p->idWrite.op_type = OP_TYPE_SEND;
		p->idWrite.pv = p;
		assert( p->send_len != 0 );
		::StartThreadpoolIo(p->pIO);
		if ( !p->sock.GetSocket().Write(p->send_buffer + p->act_sent_len, p->send_len - p->act_sent_len, NULL,
						&(p->idWrite.overlapped)) ) {
			if ( ::WSAGetLastError() != WSA_IO_PENDING ) {
				::CancelThreadpoolIo(p->pIO);
				return false;
			}
		}
		set_timer(p, false);
		return true;
	}

	void check_and_enable_timer(handle_server_data_base* p) noexcept
	{
		if ( ::InterlockedCompareExchange(&(p->lCloseFlag), 1, 1) )
			return ;
		if ( !::IsThreadpoolTimerSet(p->pTimer) )
			set_timer(p, false);
	}

	void set_timer(handle_server_data_base* p, bool bDisabled) noexcept
	{
		if ( bDisabled ) {
			::SetThreadpoolTimer(p->pTimer, NULL, 0, 0);
			return ;
		}
		LARGE_INTEGER li;
		li.QuadPart = -(int64_t)(m_util.timeout) * 1000 * 10;  //every 100 ns
		FILETIME ft;
		ft.dwLowDateTime = li.LowPart;
		ft.dwHighDateTime = li.HighPart;
		::SetThreadpoolTimer(p->pTimer, &ft, 0, 0);
	}

	static void close_handle(handle_server_data_base* p) noexcept
	{
		//first: timer
		if ( p->pTimer != NULL ) {
			::SetThreadpoolTimer(p->pTimer, NULL, 0, 0);
			::WaitForThreadpoolTimerCallbacks(p->pTimer, TRUE);
		}
		//first: cancel IO
		::CancelIoEx((HANDLE)(SOCKET)(p->sock.GetSocket()), NULL);  //no check
		p->sock.GetSocket().Disable(SD_BOTH);  //no check
		if ( p->pIO != NULL ) {
			::WaitForThreadpoolIoCallbacks(p->pIO, FALSE);
		}
		//close socket
		p->sock.GetSocket().Close();
		//second: IO
		if ( p->pIO != NULL ) {
			::WaitForThreadpoolIoCallbacks(p->pIO, TRUE);
			::CloseThreadpoolIo(p->pIO);
			p->pIO = NULL;
		}
		//second: timer
		if ( p->pTimer != NULL ) {
			::SetThreadpoolTimer(p->pTimer, NULL, 0, 0);
			::WaitForThreadpoolTimerCallbacks(p->pTimer, TRUE);
			::CloseThreadpoolTimer(p->pTimer);
			p->pTimer = NULL;
		}
	}

private:
	utilClass m_util;
};

////////////////////////////////////////////////////////////////////////////////
